import os

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.core import load_index_from_storage  # 新增加载索引功能
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope

# 配置存储路径
PERSIST_DIR = "D:/pythonRoom/llm_llama/storage/data"  # 向量索引存储目录
data_dir = 'D:/pythonRoom/llm_llama/data'
'''
阿里百炼地址
https://bailian.console.aliyun.com/?tab=app#/api-key
'''
# 1. 配置阿里云API（替换成您的API Key）
print('初始化模型对象')
llm = DashScope(
    model="qwen-turbo",  # 使用通义千问Max模型 qwen-max
    api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
    temperature=0.3
)

# 2. 使用本地嵌入模型（节省成本）
print('使用本地嵌入模型（节省成本）')
embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",  # 中文优化模型
    device="cpu"
)

# 3. 检查是否已有存储的索引
if not os.path.exists(PERSIST_DIR):
    # 3.1 加载文档
    print('加载文档')
    documents = SimpleDirectoryReader(data_dir).load_data()

    # 4. 构建索引
    print('构建索引')
    index = VectorStoreIndex.from_documents(
        documents,
        embed_model=embed_model
    )

    # 5. 保存索引到本地
    print(f'保存向量索引到: {PERSIST_DIR}')
    index.storage_context.persist(persist_dir=PERSIST_DIR)
else:
    # 6. 从本地加载已有索引
    print(f'从本地加载已有索引: {PERSIST_DIR}')
    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
    index = load_index_from_storage(
        storage_context=storage_context,
        embed_model=embed_model  # 必须使用相同的嵌入模型
    )

# 7. 创建查询引擎
print('创建查询引擎')
query_engine = index.as_query_engine(llm=llm)

# 8. 提问
print('提问')
response = query_engine.query("李大长是如何发家致富的?")
print("回答:", response)
